
import asyncio
from typing import Dict, Any, Optional
from concurrent.futures import ProcessPoolExecutor
import subprocess
import tempfile
import os


def _run_python_code_in_process(code: str, args: list, kwargs: dict) -> Optional[Dict[str, Any]]:
    """在进程中执行Python代码的函数"""
    try:
        with tempfile.NamedTemporaryFile(
                mode='w', encoding='utf-8', newline='\n', suffix='.py', delete=False
        ) as fw:
            fw.write(code)

        args_items = []
        for arg in list(args):
            if str(arg).startswith('-'):
                args_items.append(str(arg))
            elif str(arg).find('=') != -1:
                part1, part2 = str(arg).split('=', maxsplit=2)
                args_items.append(f'--{part1}={part2}')
            else:
                args_items.append(f'--{str(arg)}')

        for k, v in kwargs.items():
            args_items.append(f'--{k}="{v}"')

        string_args = " ".join(args_items)
        if string_args == '':
            cmd = f'python {fw.name}'
        else:
            cmd = f'python {fw.name} {string_args}'

        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            encoding='utf-8',
            shell=True
        )
        output, error = process.communicate()
        return {
            "success": True,
            "output": output,
            "error": error or None
        }
    except Exception as e:
        return {
            "success": False,
            "output": '',
            "error": str(e)
        }
    finally:
        if os.path.exists(fw.name):
            os.remove(fw.name)

class CodeExecutor:

    def __init__(self, timeout: int = 30, max_workers: int = 10):
        self.timeout = timeout
        self.process_pool = ProcessPoolExecutor(max_workers=max_workers)
    
    async def shutdown(self):
        """关闭进程池"""
        self.process_pool.shutdown(wait=True)

    async def execute(self, code: str, args: list, kwargs: dict) -> Dict[str, Any]:

        try:
            loop = asyncio.get_event_loop()
            future = loop.run_in_executor(
                self.process_pool,
                _run_python_code_in_process,code, args, kwargs
            )
            result = await asyncio.wait_for(future, timeout=self.timeout)
            return result

        except asyncio.TimeoutError:
            return {
                "success": False,
                "output": "",
                "error": f"代码执行超时 (>{self.timeout}秒)"
            }
        except Exception as e:
            return {
                "success": False,
                "output": "",
                "error": str(e)
            }
